/*
 * Copyright (c) 2017, AEye
 * All rights reserved.
 */

#ifdef _MSC_VER
#define _WINSOCKAPI_
#define FD_SETSIZE 1024
#include <windows.h>
#include <winsock2.h>
#include <ws2tcpip.h>
#include <stdio.h>
#define GET_LAST_ERROR_AS_STR() GetLastErrorAsString(__FILE__, __FUNCSIG__, __LINE__)
#ifndef MSG_NOSIGNAL
#define MSG_NOSIGNAL 0
#include "AEyeExceptions.hpp"
#endif
#else
#include "arpa/inet.h"
#include <unistd.h>
#include <sys/poll.h>
#define GET_LAST_ERROR_AS_STR() GetLastErrorAsString(__FILE__, __PRETTY_FUNCTION__, __LINE__)
#endif

#include <fcntl.h>
#include <cstring>
#include <iostream>
#include <chrono>
#include <sstream>
#include <system_error>

#include "spdlog/spdlog.h"

#include "NonBlockTCPSocket.h"

#ifdef _MSC_VER
//Returns the last Win32 error, in string format. Returns an empty string if there is no error.
std::string NonBlockTCPSocket::GetLastErrorAsString(std::string file, std::string function, uint16_t line)
{
    //Get the error message, if any.
    DWORD errorMessageID = ::GetLastError();
    if(errorMessageID == 0)
        return std::string(); //No error message has been recorded

    LPSTR messageBuffer = nullptr;
    size_t size = FormatMessageA(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
                                 NULL, errorMessageID, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), (LPSTR)&messageBuffer, 0, NULL);

    std::string message(messageBuffer, size);

    //Free the buffer.
    LocalFree(messageBuffer);

    return message;
}

inline
int NonBlockTCPSocket::GetLastError() {
  return ::GetLastError();
}
#else
inline
std::string NonBlockTCPSocket::GetLastErrorAsString(std::string file, std::string function, size_t line)
{
  std::cerr << strerror(errno) <<  ", errno:" << errno << "\n in file=" << file << ", " << function << ":" << line << std::endl;
  std::stringstream ss;
  ss << "Error in socket call, errno:" << errno << std::endl;
  //convert the stream buffer into a string
  return ss.str();
}

inline
int NonBlockTCPSocket::GetLastError() {
  return errno;
}
#endif

size_t NonBlockTCPSocket::send(const void* buff, size_t buffSize, size_t mSec) {
  if (!_isOpen) return 0;
  std::chrono::high_resolution_clock::time_point startTime = std::chrono::high_resolution_clock::now();
  size_t total = 0;            // how many bytes we've sent
  size_t bytesLeft = buffSize; // how many we have left to send
  ssize_t nSend;

  size_t timeGone = 0;

  while (total < buffSize && timeGone < mSec) {
    try {
      nSend = sendWithSelect(reinterpret_cast<const char*>(buff) + total, bytesLeft, mSec);

      auto endTime = std::chrono::high_resolution_clock::now();
      auto timeSpan = std::chrono::duration_cast<std::chrono::milliseconds>(endTime - startTime);
      timeGone = static_cast<size_t>(timeSpan.count());
    } catch (std::system_error& err) {
      if (err.code().value() != EAGAIN && err.code().value() != EINTR) {
        throw;
      }
      auto endTime = std::chrono::high_resolution_clock::now();
      auto timeSpan = std::chrono::duration_cast<std::chrono::milliseconds>(endTime - startTime);
      timeGone = static_cast<size_t>(timeSpan.count());
      continue;
    }
    if (nSend < 0) {// if (nSend == -1) {
      break;
    }
    total     += nSend;
    bytesLeft -= nSend;
  }
  return total;
}

size_t NonBlockTCPSocket::receive(void* buff, size_t buffSize, size_t mSec) {
  if (!_isOpen) return 0;
  auto startTime = std::chrono::high_resolution_clock::now();
  size_t total = 0;            // how many bytes we've sent
  size_t bytesLeft = buffSize; // how many we have left to send
  ssize_t nRecv;

  size_t timeGone = 0;

  while (total < buffSize && timeGone < mSec) {
    try {
      nRecv = receiveWithSelect(reinterpret_cast<char*>(buff) + total, bytesLeft, mSec);

      auto endTime = std::chrono::high_resolution_clock::now();
      auto timeSpan = std::chrono::duration_cast<std::chrono::milliseconds>(endTime - startTime);
      timeGone = static_cast<size_t>(timeSpan.count());
    } catch (std::system_error& err) {
      auto e = err.code().value();
      if (e == EHOSTUNREACH) {//113
        spdlog::warn("Host unreachable.");
        close();
      }

      if (e != EAGAIN && e != EINTR) {
        spdlog::error("Received error code \"{}\" in NonBlockTCPSocket::receive.", e);
        throw;
      } else {
        spdlog::warn("Received error code \"{}\" in NonBlockTCPSocket::receive.", e);
      }
      auto endTime = std::chrono::high_resolution_clock::now();
      auto timeSpan = std::chrono::duration_cast<std::chrono::milliseconds>(endTime - startTime);
      timeGone = static_cast<size_t>(timeSpan.count());
      continue;
    }
    if (nRecv < 0) {
      break;
    }
    total     += nRecv;
    bytesLeft -= nRecv;
  }
  return total;
}

ssize_t NonBlockTCPSocket::sendWithSelect(const void* buff, size_t buffSize, size_t mSec) {

  //assert(fd > 0 && fd <= FD_SETSIZE);
  struct timeval tv;
  tv.tv_sec = mSec / 1000;
  tv.tv_usec = 1000 * (mSec - tv.tv_sec * 1000);

  fd_set writeFlags;
  FD_ZERO(&writeFlags);
  FD_SET(fd, &writeFlags);

  // don't care about readfds and exceptfds:
  int selectReturn = select(fd + 1, nullptr, &writeFlags, nullptr, &tv);
  if (selectReturn == -1) {
    // error happened
    int e = errno;
    spdlog::warn("UDP select reports `{1}`, errno {0}.", e, strerror(e));
    throw std::system_error(GetLastError(), std::system_category(), strerror(e));
  }
  else if (selectReturn == 0) {
    // mSec
    return 0;
  }

  if (FD_ISSET(fd, &writeFlags)) {
    // ready to send
    ssize_t nSend = ::send(fd, reinterpret_cast<const char*>(buff), buffSize, MSG_NOSIGNAL);
    if (nSend == -1) {
      auto e = errno;
      spdlog::error("UDP send reports `{1}`, errno {0}.", e, strerror(e));
      throw std::system_error(e, std::system_category(), strerror(e));
    }
    return nSend;
  }

  return 0;
}

ssize_t NonBlockTCPSocket::receiveWithSelect(void* buff, size_t buffSize, size_t mSec) {

  //assert(fd > 0 && fd <= FD_SETSIZE);
  if (getSO_ERROR() != 0) {
    throw std::system_error(GetLastError(), std::system_category(), GET_LAST_ERROR_AS_STR());
  }

  struct timeval tv;
  tv.tv_sec = mSec / 1000;
  tv.tv_usec = 1000 * (mSec - tv.tv_sec * 1000);

  fd_set readFlags;
  FD_ZERO(&readFlags);
  FD_SET(fd, &readFlags);

  // don't care about writefds and exceptfds:
  int selectReturn = select(fd + 1, &readFlags, nullptr, nullptr, &tv);
  if (selectReturn == -1) {
    // error happened
    auto e = errno;
    spdlog::warn("UDP select reports `{1}`, errno {0}.", e, strerror(e));
    throw std::system_error(e, std::system_category(), strerror(e));
  }
  else if (selectReturn == 0) {
    // timeout
    return 0;
  }

  if (FD_ISSET(fd, &readFlags)) {
    // ready to read
    ssize_t nReceive = ::recv(fd, reinterpret_cast<char*>(buff), buffSize, MSG_NOSIGNAL); // receive normal data
    if (nReceive == -1) {
      auto e = errno;
      spdlog::warn("UDP recv reports `{1}`, errno {0}.", e, strerror(e));
      throw std::system_error(e, std::system_category(), strerror(e));
    }
    else if (nReceive == 0) {
      // connection was closed by the server
    }
    return nReceive;
  } // else { // Timeout or wrong checksum }
  return 0;
}


bool NonBlockTCPSocket::isConnected() {
  if (!_isOpen) return false;
#ifdef WIN32
  char buf;
#ifdef _MSC_VER
  int err = ::recv(fd, &buf, 1, MSG_PEEK | 0);
#else
  int err = ::recv(fd, &buf, 1, MSG_PEEK | MSG_NOSIGNAL);
#endif
  if (err == SOCKET_ERROR) {
    if (WSAGetLastError() != WSAEWOULDBLOCK) {
      return false;
    }
  }
  else if (err == 0) {
    return false; // Server closed connection
  }
  return true;
#else
  int res = 0;
  try {
    res = getSO_ERROR();
  } catch (std::system_error& err) {
    std::cout << "Error: " << err.code() << " - " << err.what() << '\n';
    return false;
  }
  if (res == 0) {
    int pollReturn;
    struct pollfd ufds;
    // set up the array of file descriptors.
    ufds.fd = fd;
    ufds.events = POLLIN; // check for normal data in socket.

    // wait for events on the sockets, with timeout
    pollReturn = poll(&ufds, 1, 0);
    if (pollReturn == -1) {
      perror("poll"); // error occurred in poll()
      return false;
    } else if (pollReturn == 0) {
      //printf("Timeout occurred! \n");
      return true;
    } else {
      // check for events on socket:
      if (ufds.revents & (POLLERR | POLLHUP | POLLNVAL)) {
        // An error has occurred on this socket.
        return false;
      }
      char data;
      ssize_t nrecv = ::recv(fd, &data, 1, MSG_PEEK | MSG_NOSIGNAL); //peak one byte
      if (nrecv == -1 || nrecv == 0) {
        return false;
      }
      else return true;
    }
  }
#endif
  return false;
}

void NonBlockTCPSocket::cleanUp(size_t mSec) {
  if (getSO_ERROR() != 0) {
    throw std::system_error(GetLastError(), std::system_category(), GET_LAST_ERROR_AS_STR());
  }

  struct timeval tv;
  tv.tv_sec = mSec / 1000;
  tv.tv_usec = 1000 * (mSec - tv.tv_sec * 1000);

  fd_set readFlags;
  FD_ZERO(&readFlags); // Initialize to zero's
  FD_SET(fd, &readFlags);

  // don't care about writefds and exceptfds:
  int selectReturn = select(fd + 1, &readFlags, nullptr, nullptr, &tv);
  if (selectReturn == -1) {
    // error happened
    int e = errno;
    spdlog::warn("UDP select reports `{1}`, errno {0}.", e, strerror(e));
    throw std::system_error(GetLastError(), std::system_category(), strerror(e));
  }
  else if (selectReturn == 0) {
    // timeout
    return;
  }

  if (FD_ISSET(fd, &readFlags)) {
    // ready to read
    char data;
    ssize_t nBytesAvailable = ::recv(fd, &data, 1, MSG_PEEK | MSG_NOSIGNAL); //peak at one byte
    if (nBytesAvailable == -1) {
      int e = errno;
      spdlog::warn("UDP recv reports `{1}`, errno {0}.", e, strerror(e));
      throw std::system_error(GetLastError(), std::system_category(), strerror(e));
    }
    else if (nBytesAvailable > 0) {
      char buf[4096];
      ssize_t bytesRead;
      do {
        bytesRead = ::recv(fd, buf, sizeof(buf), MSG_NOSIGNAL);
        // Do nothing just discard data from socket
      } while (bytesRead > 0);
    }
  } // else { // Timeout or wrong checksum }
}

int NonBlockTCPSocket::setNonblockingAttribute() {
  int flags = 0;
  /* If we have O_NONBLOCK, use the Posix way to do it */
#if defined(O_NONBLOCK)
  /* Fixme: O_NONBLOCK is defined but broken on SunOS 4.1.x and AIX 3.2.5. */
  flags = fcntl(fd, F_GETFL, 0);
  if (flags == -1)
    return -1;

  // flags &= (~O_NONBLOCK);  // for blocking
  // flags |= O_NONBLOCK;     // for non-blocking
  return fcntl(fd, F_SETFL, flags | O_NONBLOCK);
#else
  /* Otherwise, use the old way of doing it */
#ifdef _MSC_VER
  
  // Creating non blocking sockets on windows is quite difficult, especially when this is built around unix sockets. 
  // Sadly, I'm not a socket expert so using this method our recieves will still block but will time out quite quickly
  // thus they are essentially non blocking. I set the timout to 100 to have a small grace period if the off chance it can't
  // connect immediately.
  int timeoutThreshold = 100;
  if (setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, reinterpret_cast<char*>(&timeoutThreshold), sizeof(timeoutThreshold)) != 0) {
      spdlog::error("setsockopt reports `{1}`, errno {0}.", errno, strerror(errno));
      throw aeye::socket_error();
  }

  flags = 0; // The interpretation of the flags argument differs between WinSock and UNIX sockets
  return ioctlsocket(fd, FIONBIO, (u_long*)&flags);
#else
  flags = 1;
  return ioctl(fd, FIOBIO, &flags);
#endif
#endif
}

NonBlockTCPSocket::NonBlockTCPSocket() : _isOpen(false), fd(-1) {
#if 0 // need to call these on upper level only once
#if _WIN32 || _WIN64
    WORD wVersionRequested;
    WSADATA wsaData;
    int err;

/* Use the MAKEWORD(lowbyte, highbyte) macro declared in Windef.h */
    wVersionRequested = MAKEWORD(2, 2);

    err = WSAStartup(wVersionRequested, &wsaData);
    if (err != 0) {
        /* Tell the user that we could not find a usable */
        /* Winsock DLL.                                  */
        printf("WSAStartup failed with error: %d\n", err);
        throw std::system_error(err, std::system_category(), "WSAStartup failed.");
    }

/* Confirm that the WinSock DLL supports 2.2.*/
/* Note that if the DLL supports versions greater    */
/* than 2.2 in addition to 2.2, it will still return */
/* 2.2 in wVersion since that is the version we      */
/* requested.                                        */

    if (LOBYTE(wsaData.wVersion) != 2 || HIBYTE(wsaData.wVersion) != 2) {
        /* Tell the user that we could not find a usable */
        /* WinSock DLL.                                  */
        printf("Could not find a usable version of Winsock.dll\n");
        WSACleanup();
        throw std::system_error(-1, std::system_category(), "Could not find a usable version of Winsock.dll.");
    }
    else
        printf("The Winsock 2.2 dll was found okay\n");
#endif
#endif // #if 0
}

NonBlockTCPSocket::~NonBlockTCPSocket() {
  spdlog::trace("Closing NonBlockTCPSocket.");
  close();
}

bool NonBlockTCPSocket::isOpen() const {
  return _isOpen;
}

int NonBlockTCPSocket::getSO_ERROR() {
  int err = 1;
  socklen_t len = sizeof(err);
  if (-1 == getsockopt(fd, SOL_SOCKET, SO_ERROR, reinterpret_cast<char *>(&err), &len)) {
    spdlog::warn("getsockopt reports `{1}`, errno {0}.", errno, strerror(errno));
    return -1;
  }
#if  !_WIN32 && !_WIN64
  if (err)
    errno = err;              // set errno to the socket SO_ERROR
#endif
  return err;
}

void NonBlockTCPSocket::close() {
  if (fd >= 0) {
    getSO_ERROR(); // first clear any errors, which can cause close to fail
#ifdef _MSC_VER
  shutdown(fd, SD_BOTH);
  if (::closesocket(fd) != 0) {
    if (GetLastError() != WSAENOTSOCK) { // if any error among "The descriptor is not a socket".
      throw std::system_error(GetLastError(), std::system_category(), GET_LAST_ERROR_AS_STR());
    }
  }
  // WSACleanup();
#else
    if (shutdown(fd, SHUT_RDWR) < 0) // secondly, terminate the 'reliable' delivery
      if (errno != ENOTCONN && errno != EINVAL) {// SGI causes EINVAL
        spdlog::warn("UDP shutdown reports `{1]`, errno {0}.", errno, strerror(errno));
        getSO_ERROR(); // clear any errors, which can cause close to fail
      }
    if (::close(fd) < 0) { // finally call close()
      spdlog::warn("UDP close reports `{1}`, errno {0}.", errno, strerror(errno));
      if (errno == EINTR) { // The close() function was interrupted by a signal.
        // Try again
        if (::close(fd) < 0 && errno == EINTR) {// The close() function was interrupted by a signal.
          spdlog::warn("UDP close (attempt 2) reports `{1}`, errno {0}.", errno, strerror(errno));
          throw std::system_error(GetLastError(), std::system_category(), GET_LAST_ERROR_AS_STR());
        }
        _isOpen = false;
        fd = -1;
      }
      else if (errno == EIO) { // An I/O error occurred.
        throw std::system_error(GetLastError(), std::system_category(), GET_LAST_ERROR_AS_STR());
      } // else { // skip other errors if socket has been closed in another operations
    }
#endif
    _isOpen = false;
    fd = -1;
  }
}

bool NonBlockTCPSocket::connect(const std::string& addr, uint16_t port, size_t timeoutSec) {

  if (_isOpen) {
    try {
      this->close();
    }
    catch (std::exception& e) {
      spdlog::warn("UDP socket close raised exception: {}.", e.what());
    }
  }

  // Create socket
  fd = socket(AF_INET, SOCK_STREAM, 0);
  if (fd < 0) {
    spdlog::error("UDP socket ctor reports `{1}`, errno {0}.", errno, strerror(errno));
    throw std::system_error(GetLastError(), std::system_category(), GET_LAST_ERROR_AS_STR());
  }
  // setup non blocking socket
  if (setNonblockingAttribute() != 0) {
    spdlog::error("Unable to set nonblocking attribute, reports `{1}`, errno {0}.", errno, strerror(errno));
    int errnoPrev = errno;
    getSO_ERROR(); // clear any errors, which can cause close to fail
#if _WIN32 || _WIN64
    closesocket(fd);
#else
    ::close(fd);
#endif
    fd = -1;
    errno = errnoPrev;
    throw std::system_error(GetLastError(), std::system_category(), GET_LAST_ERROR_AS_STR());
  }
#ifdef _MSC_VER
  // Windows already checks if fd is valid and will return fd < 0, if there is an error (Output WSAGetLastError() to find Windows error code)
  if(fd < 0) {
#else
  if (fd < 0 || fd > FD_SETSIZE) { // this check needs for properly operation with select method
#endif
    spdlog::error("nonblock socket connect failed, FD out of range.");
    getSO_ERROR(); // clear any errors, which can cause close to fail
#if _WIN32 || _WIN64
    closesocket(fd);
#else
    ::close(fd);
#endif
    fd = -1;
    errno = EBADF;
    throw std::system_error(errno, std::system_category(), GET_LAST_ERROR_AS_STR());
  }

  struct sockaddr_in servAddr = sockaddr_in(); // name initializer to silence spurious gcc warning
  // Set up the socket options
  servAddr.sin_family       = AF_INET;
  servAddr.sin_addr.s_addr  = inet_addr(addr.c_str());
  servAddr.sin_port         = htons(port);

  // make the connection
  int rc = ::connect(fd, reinterpret_cast<struct sockaddr*>(&servAddr), sizeof(servAddr));

#ifdef _MSC_VER
  if ((rc == -1) && (GetLastError() != WSAEINPROGRESS)) {
    std::cerr << "Error NonBlockTCPSocket()::connect :" << GET_LAST_ERROR_AS_STR() << std::endl;
   ::closesocket(fd);
#else
  if ((rc == -1) && (GetLastError() != EINPROGRESS)) {
    spdlog::error("UDP connect reports `{1}`, errno {0}. (rc = {2}).", errno, strerror(errno), rc);
    int errnoPrev = errno;
    getSO_ERROR(); // clear any errors, which can cause close to fail
    ::close(fd);
    fd = -1;
    errno = errnoPrev;
#endif
    throw std::system_error(GetLastError(), std::system_category(), GET_LAST_ERROR_AS_STR());
  }
  if (rc == 0) {
    // connection has succeeded immediately
    _isOpen = true;
    return true;
  } // else {
        // connection attempt is in progress }

  fd_set writefds;
  FD_ZERO(&writefds);
  FD_SET(fd, &writefds);
  struct timeval tv{};
  tv.tv_sec = timeoutSec;
  tv.tv_usec = 0;

  rc = select(fd + 1, nullptr, &writefds, nullptr, &tv);
#if _WIN32 || _WIN64
  if (rc == 1) {
    _isOpen = true;
    return true; // Connected !!
  }
  else if (rc == 0) { // Timeout !!
    return false;
  }
  else { // Connection error
    throw std::system_error(GetLastError(), std::system_category(), GET_LAST_ERROR_AS_STR());
  }
#else
  if (rc == 1) {
    if (getSO_ERROR() == 0) {  // Connected !!!
      _isOpen = true;
      return true;
    }
    else { // error
      ::close(fd);
      fd = -1;
      throw std::system_error(GetLastError(), std::system_category(), GET_LAST_ERROR_AS_STR());
    }
  }
  else if (rc == 0) { // Timeout !!
    getSO_ERROR(); // clear any errors, which can cause close to fail
    ::close(fd);
    fd = -1;
    return false;
  }
  else { // Connection error
    int errnoPrev = errno;
    getSO_ERROR(); // clear any errors, which can cause close to fail
    ::close(fd);
    fd = -1;
    errno = errnoPrev;
    throw std::system_error(GetLastError(), std::system_category(), GET_LAST_ERROR_AS_STR());
  }
#endif
}
